package at.jku.cp.mir;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeMap;

import org.apache.commons.io.IOUtils;
import org.jdom.Document;
import org.jdom.Element;
import org.jdom.JDOMException;
import org.jdom.input.SAXBuilder;
import org.jdom.xpath.XPath;
import org.json.simple.JSONArray;
import org.json.simple.JSONObject;
import org.json.simple.JSONValue;

import com.google.common.collect.TreeMultimap;

/**
 * This class reads handles the calculation of similar artists.
 * @author clemens
 */
public class SimilarityCalculator
{

	private ArtistVectorSimilarity similarity;

	public SimilarityCalculator()
	{
		super();

		this.similarity = new ArtistVectorCosineSimilarity();
	}

	/**
	 * Read the tags from the given directories, calculate the similar artists
	 * and store the file in <code>similarityFile</code>.
	 * @param artistTopTagsDir
	 * @param topTrackTagsDir
	 * @param similarityFile
	 * @throws JDOMException
	 * @throws IOException
	 */
	public void createSimilarArtistsFile(File artistTopTagsDir, File topTrackTagsDir, File similarityFile)
			throws JDOMException, IOException
	{
		long start = System.currentTimeMillis();

		// create the artist tag maps by reading the entries in the directory
		Map<String, List<Tag>> artistTopTagMap = readArtistTopTagMapFromDir(artistTopTagsDir);
		Map<String, List<Tag>> artistTopTrackTagsMap = readArtistTopTrackTagsMapFromDir(topTrackTagsDir);

		long end = System.currentTimeMillis();
		System.out.println("it took " + (end - start) + "ms to create the maps");

		// finally create the similarity artists file
		createSimilarArtistsFile(artistTopTagMap, artistTopTrackTagsMap, similarityFile);
	}

	public void createSimilarArtistsFile(Map<String, List<Tag>> artistTopTagMap,
			Map<String, List<Tag>> artistTopTrackTagsMap, File similarityFile)
	{
		long start = System.currentTimeMillis();

		// normalize the tags
		TagNormalizer tagNormalizer = new TagNormalizer();
		Map<String, List<Tag>> normalizedTags = tagNormalizer.normalizeTags(artistTopTagMap);
		normalizedTags = tagNormalizer.applyTrackFiltering(normalizedTags, artistTopTrackTagsMap);

		// calculate TfIdf'S
		TfIdfCalculator tfIdfCalculator = new TfIdfCalculator();
		Map<String, ArtistVector> artistVectors = tfIdfCalculator.calculateTfIdf(normalizedTags);

		// calculate similar artists
		Map<String, List<String>> similarArtists = calculateSimialrArtists(artistVectors);

		long end = System.currentTimeMillis();
		System.out.println("it took " + (end - start) + "ms to calculate simiar artists");

		// output similar artists
		Set<String> artists = similarArtists.keySet();
		List<String> lines = new ArrayList<String>(artists.size());
		for (String artist : artists)
		{
			List<String> similar = similarArtists.get(artist);
			StringBuilder similarArtistsString = new StringBuilder(artist);
			similarArtistsString.append(" --> ");
			if (similar.size() > 0)
			{
				for (String similarArtist : similar)
				{
					similarArtistsString.append(similarArtist).append(",");
				}
				// delete last char
				similarArtistsString.reverse().delete(0, 1).reverse();
			}
			lines.add(similarArtistsString.toString());
		}
		try
		{
			IOUtils.writeLines(lines, "\n", new FileOutputStream(similarityFile));
		} catch (FileNotFoundException e)
		{
			e.printStackTrace();
		} catch (IOException e)
		{
			e.printStackTrace();
		}
	}

	public Map<String, List<Tag>> readArtistTopTagMapFromDir(File artistTopTagsDir) throws JDOMException, IOException
	{
		Map<String, List<Tag>> artistTopTagMap = new TreeMap<String, List<Tag>>();

		SAXBuilder xmlBuilder = new SAXBuilder();
		Document artistTopTagsDoc = null;

		// get all artists and their tags
		File[] artistTopTagsFiles = artistTopTagsDir.listFiles();
		for (File artistTopTagsFile : artistTopTagsFiles)
		{
			artistTopTagsDoc = xmlBuilder.build(artistTopTagsFile);

			String artistName = artistTopTagsFile.getName();
			artistName = artistName.replace(".xml", "");
			artistName = artistName.replace("_", "/");
			List<Tag> topTags = extractTopTags(artistTopTagsDoc);

			artistTopTagMap.put(artistName, topTags);
		}

		return artistTopTagMap;
	}

	public Map<String, List<Tag>> readArtistTopTagMapFromFile(File artistTopTagsFile) throws FileNotFoundException
	{
		return readMapFromFile(artistTopTagsFile);
	}

	public Map<String, List<Tag>> readArtistTopTrackTagsMapFromDir(File topTrackTagsDir) throws JDOMException,
			IOException
	{
		Map<String, List<Tag>> artistTopTrackTagsMap = new TreeMap<String, List<Tag>>();

		SAXBuilder xmlBuilder = new SAXBuilder();

		// get all top tags for the tracks of each artist
		File[] artistDirs = topTrackTagsDir.listFiles();
		for (File artistDir : artistDirs)
		{
			File[] topTrackTagFiles = artistDir.listFiles();
			String artistName = artistDir.getName();
			artistName = artistName.replace("_", "/");
			List<Tag> topTrackTags = new LinkedList<Tag>();
			for (File topTrackTagFile : topTrackTagFiles)
			{
				Document topTrackTagDoc = xmlBuilder.build(topTrackTagFile);
				List<Tag> tmpTagList = extractTopTags(topTrackTagDoc);
				topTrackTags.addAll(tmpTagList);
			}

			artistTopTrackTagsMap.put(artistName, topTrackTags);
		}

		return artistTopTrackTagsMap;
	}

	public Map<String, List<Tag>> readArtistTopTrackTagsMapFromFile(File topTrackTagsFile) throws FileNotFoundException
	{
		return readMapFromFile(topTrackTagsFile);
	}

	public void setArtistVectoSimilarity(ArtistVectorSimilarity similarity)
	{
		this.similarity = similarity;
	}

	public void writeArtistTopTagMapToFile(Map<String, List<Tag>> artistTopTagMap, File artistTopTagMapFile)
			throws FileNotFoundException, IOException
	{
		JSONObject jsonObject = toJSONObject(artistTopTagMap, "artistTopTagMap");
		IOUtils.write(jsonObject.toJSONString(), new FileOutputStream(artistTopTagMapFile));
	}

	public void writeArtistTopTrackTagMapToFile(Map<String, List<Tag>> artistTopTrackTags, File artistTopTrackTagMapFile)
			throws FileNotFoundException, IOException
	{
		JSONObject jsonObject = toJSONObject(artistTopTrackTags, "artistTopTrackTagMap");
		IOUtils.write(jsonObject.toJSONString(), new FileOutputStream(artistTopTrackTagMapFile));
	}

	private Map<String, List<String>> calculateSimialrArtists(Map<String, ArtistVector> artistVectors)
	{
		Set<String> artists = artistVectors.keySet();
		Iterator<String> artistIt = artists.iterator();

		Map<String, List<String>> artistSimilarityMap = new TreeMap<String, List<String>>();
		while (artistIt.hasNext())
		{
			String currentArtist = artistIt.next();
			ArtistVector firstArtistVector = artistVectors.get(currentArtist);

			Iterator<String> secondArtistIt = artists.iterator();
			TreeMultimap<Double, String> similarityMap = TreeMultimap.create();
			while (secondArtistIt.hasNext())
			{
				String secondArtist = secondArtistIt.next();
				if (!currentArtist.equals(secondArtist))
				{
					ArtistVector av = artistVectors.get(secondArtist);
					double similarity = this.similarity.getSimilarity(firstArtistVector, av);
					similarityMap.put(Double.valueOf(similarity), secondArtist);
				}
			}

			LinkedList<String> linkedSimilarArtists = new LinkedList<String>();
			SortedSet<Double> keySet = similarityMap.keySet();
			Iterator<Double> iterator = keySet.iterator();
			while (iterator.hasNext())
			{
				Double next = iterator.next();
				SortedSet<String> nextSet = similarityMap.get(next);
				for (String nextSimilarArtist : nextSet)
				{
					linkedSimilarArtists.addFirst(nextSimilarArtist);
				}
			}

			Double similarity = keySet.last();
			SortedSet<String> similarArtists = similarityMap.get(similarity);
			// System.out.println(currentArtist + " --> similar artists (by " +
			// similarity + "): "
			// + Arrays.toString(similarArtists.toArray()));

			// artistSimilarityMap.put(currentArtist, new
			// ArrayList<String>(similarArtists));
			int index = Math.min(3, linkedSimilarArtists.size());
			index = linkedSimilarArtists.size();
			artistSimilarityMap.put(currentArtist, linkedSimilarArtists.subList(0, index));
		}
		return artistSimilarityMap;
	}

	private List<Tag> extractTopTags(Document doc) throws JDOMException
	{
		List<Tag> tagList = new LinkedList<Tag>();
		List<Element> tags = XPath.selectNodes(doc, "//tag");
		for (Element tagElement : tags)
		{
			String name = tagElement.getChild("name").getTextTrim();
			String countString = tagElement.getChild("count").getTextTrim();
			int count = Integer.parseInt(countString);

			tagList.add(new Tag(name, count));
		}
		return tagList;
	}

	private Map<String, List<Tag>> readMapFromFile(File artistTopTagsFile) throws FileNotFoundException
	{
		JSONObject jsonObj = (JSONObject) JSONValue.parse(new FileReader(artistTopTagsFile));

		Map<String, List<Tag>> map = new HashMap<String, List<Tag>>();
		JSONArray artistsArr = (JSONArray) jsonObj.get("artists");
		for (Object object : artistsArr)
		{
			JSONObject obj = (JSONObject) object;
			String artist = (String) obj.get("artistName");
			JSONArray tags = (JSONArray) obj.get("tags");
			List<Tag> tagList = new ArrayList<Tag>(tags.size());
			for (Object tmp : tags)
			{
				JSONObject tagObj = (JSONObject) tmp;
				String name = (String) tagObj.get("name");
				int count = ((Number) tagObj.get("count")).intValue();
				Tag tag = new Tag(name, count);
				tagList.add(tag);
			}
			map.put(artist, tagList);
		}

		return map;
	}

	private JSONObject toJSONObject(Map<String, List<Tag>> map, String typeOfContent)
	{
		JSONObject mapFile = new JSONObject();
		mapFile.put("content", typeOfContent);
		JSONArray artistsArray = new JSONArray();
		Set<String> artists = map.keySet();
		for (String artist : artists)
		{
			List<Tag> tags = map.get(artist);
			JSONObject artistObject = new JSONObject();
			artistObject.put("artistName", artist);
			artistObject.put("tags", tags);
			artistsArray.add(artistObject);
		}
		mapFile.put("artists", artistsArray);
		return mapFile;
	}
}
